package com.leetcode;

import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * 347. 前 K 个高频元素
 * 通过堆排序 时间复杂度O(nlogk)
 * 但是内部用数组同时存数据及数据堆个数
 */
public class Solution347_1 {

    public int[] topKFrequent(int[] nums, int k) {
        final Map<Integer, Integer> map = new HashMap<>();
        for (int num : nums) {
            map.put(num, map.getOrDefault(num, 0) + 1);
        }
        // java默认的堆是最小堆，这里通过数组同时存数据和数据个数
        PriorityQueue<int[]> priorityQueue = new PriorityQueue<>(Comparator.comparingInt((int[] a) -> a[1]));
        for (Integer num : map.keySet()) {
            if (priorityQueue.size() >= k) {
                if (map.get(num) > priorityQueue.peek()[1]) {
                    priorityQueue.poll();
                    priorityQueue.offer(new int[]{num, map.get(num)});
                }
            } else {
                priorityQueue.offer(new int[]{num, map.get(num)});
            }
        }
        int[] res = new int[k];
        int index = 0;
        while (!priorityQueue.isEmpty()) {
            res[index] = priorityQueue.poll()[0];
            index++;
        }
        return res;
    }

    public static void main(String[] args) {
        int[] nums = new int[]{1, 3, 1, 2, 1, 2, 2, 3, 3, 3, 1};
        int k = 2;
        int[] res = new Solution347().topKFrequent(nums, k);
        for (int num : res) {
            System.out.println(num + " ");
        }
        System.out.println();
    }

}
